package com.hyk.learning.loadbalance.loadbalancer;

import com.hyk.learning.common.factory.SingletonFactory;
import com.hyk.learning.common.utils.AdaptiveMetrics;
import com.hyk.learning.dto.RpcRequest;
import com.hyk.learning.loadbalance.AbstractLoadBalance;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ThreadLocalRandom;

/**
 * @className: AdaptiveLoadBalance
 * @author: hyk199710
 * @description: 自适应负载均衡算法，参考Dubbo
 * @date: 2024/8/13 23:38
 */
public class AdaptiveLoadBalance extends AbstractLoadBalance {

    private static final int DEFAULT_WIEIGHT = 100;
    private static final long DEFAULT_WARMUP = 10 * 60 * 1000;
    private static final int DEFAULT_TIMEOUT = 1000;
    private Map<String, Integer> weightMap = new ConcurrentHashMap<>();
    private Map<String, Long> startTSMap = new ConcurrentHashMap<>();
    private Map<String, Long> warmupTSMap = new ConcurrentHashMap<>();
    private Map<String, Integer> timeoutMap = new ConcurrentHashMap<>();

    private final AdaptiveMetrics adaptiveMetrics;

    public AdaptiveLoadBalance() {
        adaptiveMetrics = SingletonFactory.getInstance(AdaptiveMetrics.class);
    }

    @Override
    protected String doSelect(List<String> serviceUrlList, RpcRequest rpcRequest) {
        String invoker = selectByP2C(serviceUrlList, rpcRequest);
        long startTime = System.currentTimeMillis();
        adaptiveMetrics.addConsumerReq(invoker);
        adaptiveMetrics.setPickTime(invoker, startTime);
        return invoker;
    }

    private String selectByP2C(List<String> serviceUrlList, RpcRequest rpcRequest) {
        int length = serviceUrlList.size();

        if(length == 1){
            return serviceUrlList.get(0);
        }

        if(length == 2){
            return chooseLowLoadInvoker(serviceUrlList.get(0), serviceUrlList.get(1), rpcRequest);
        }

        int pos1 = ThreadLocalRandom.current().nextInt(length);
        int pos2 = ThreadLocalRandom.current().nextInt(length - 1);
        if(pos2 >= pos1){
            pos2 = pos2 + 1;
        }

        return chooseLowLoadInvoker(serviceUrlList.get(pos1), serviceUrlList.get(pos2), rpcRequest);
    }

    private String chooseLowLoadInvoker(String invoker1, String invoker2, RpcRequest rpcRequest){
        int weight1 = getWeight(invoker1, rpcRequest);
        int weight2 = getWeight(invoker2, rpcRequest);

        if(!timeoutMap.containsKey(invoker1)){
            timeoutMap.put(invoker1, rpcRequest.getTimeout());
        }
        if(!timeoutMap.containsKey(invoker2)){
            timeoutMap.put(invoker2, rpcRequest.getTimeout());
        }
        int timeout1 = timeoutMap.get(invoker1);
        int timeout2 = timeoutMap.get(invoker2);
        long load1 = Double.doubleToLongBits(
                adaptiveMetrics.getLoad(invoker1, weight1, timeout1));
        long load2 = Double.doubleToLongBits(
                adaptiveMetrics.getLoad(invoker2, weight2, timeout2));

        if(load1 == load2){
            int totalWeight = weight1 + weight2;
            if(totalWeight > 0){
                int offset = ThreadLocalRandom.current().nextInt(totalWeight);
                if(offset > weight1){
                    return invoker1;
                }
                return invoker2;
            }
            return ThreadLocalRandom.current().nextInt(2) == 0 ? invoker1 : invoker2;
        }

        return load1 > load2 ? invoker2 : invoker1;
    }

    private int getWeight(String invoker, RpcRequest rpcRequest){
        int weight;

        if(!weightMap.containsKey(invoker)){
            weightMap.put(invoker, rpcRequest.getWeight());
        }
        weight = weightMap.get(invoker);
        if(weight > 0){
            if(!startTSMap.containsKey(invoker)){
                startTSMap.put(invoker, rpcRequest.getStartTime());
            }
            long timestamp = startTSMap.get(invoker);
            if(timestamp > 0L){
                long uptime = System.currentTimeMillis() - timestamp;
                if(uptime < 0){
                    return 1;
                }
                if(!warmupTSMap.containsKey(invoker)){
                    warmupTSMap.put(invoker, rpcRequest.getWarmupTS());
                }
                long warmup = warmupTSMap.get(invoker);
                if(uptime > 0 && uptime < warmup) {
                    int ww = (int) (uptime / ((float)warmup / weight));
                    weight = ww < 1 ? 1 : (Math.min(ww, weight));
                }
            }
        }

        return Math.max(weight, 0);
    }

}
